import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms
import os
import cv2
from sklearn.model_selection import train_test_split
import numpy as np
from torch.utils.data import Dataset, DataLoader
from PIL import Image

# 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 自定义数据集类
class CatDogDataset(Dataset):
    """
    自定义猫狗数据集类
    """
    def __init__(self, image_paths, labels, transform=None):
        """
        初始化数据集
        参数:
            image_paths: 图像路径列表
            labels: 标签列表
            transform: 数据转换操作
        """
        self.image_paths = image_paths
        self.labels = labels
        self.transform = transform

    def __len__(self):
        """
        返回数据集大小
        """
        return len(self.image_paths)

    def __getitem__(self, idx):
        """
        获取指定索引的数据
        参数:
            idx: 索引
        返回:
            图像和标签
        """
        # 读取图像
        image = cv2.imread(self.image_paths[idx])
        # 转换为RGB格式
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 转换为PIL图像
        image = Image.fromarray(image)
        # 获取标签
        label = self.labels[idx]
        
        # 应用转换：数据预处理和增强
        if self.transform:
            image = self.transform(image)
            
        return image, label

# 3. 划分训练集/验证集（80%/20%）
def prepare_data(data_dir='Data/dogs-vs-cats'):
    """
    准备训练和验证数据
    参数:
        data_dir: 数据目录路径
    返回:
        训练和验证数据的路径和标签
    """
    # 获取所有训练图像路径
    train_dir = os.path.join(data_dir, 'train')
    image_files = os.listdir(train_dir)
    
    # 创建图像路径和标签列表
    image_paths = []
    labels = []
    
    for file in image_files:
        image_paths.append(os.path.join(train_dir, file))
        # 根据文件名确定标签（cat=0, dog=1）
        if file.startswith('cat'):
            labels.append(0)
        else:  # dog
            labels.append(1)
    
    # 划分训练集和验证集 (80%/20%)
    train_paths, val_paths, train_labels, val_labels = train_test_split(
        image_paths, labels, test_size=0.2, random_state=42, stratify=labels
    )
    
    return train_paths, val_paths, train_labels, val_labels

# 定义数据预处理和增强
# transforms.Compose用于组合多个图像变换操作
data_transforms = {
    # 训练集的预处理
    'train': transforms.Compose([
        # 调整图像大小为224x224
        transforms.Resize((224, 224)),
        # 随机水平翻转（数据增强）
        transforms.RandomHorizontalFlip(),
        # 转换为tensor格式
        transforms.ToTensor(),
        # 标准化处理，均值和标准差都是基于ImageNet数据集计算的
        #项目使用的ResNet18模型是在ImageNet数据集上预训练的。
        #为了保持一致性并充分利用预训练模型的特征提取能力，需要使用相同的标准化参数。
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    # 验证集的预处理（通常不需要数据增强）
    'val': transforms.Compose([
        # 调整图像大小为224x224
        transforms.Resize((224, 224)),
        # 转换为tensor格式
        transforms.ToTensor(),
        # 标准化处理
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

# 准备数据
train_paths, val_paths, train_labels, val_labels = prepare_data()

# 创建数据集
train_dataset = CatDogDataset(train_paths, train_labels, transform=data_transforms['train'])
val_dataset = CatDogDataset(val_paths, val_labels, transform=data_transforms['val'])

# 创建数据加载器
# DataLoader用于批量加载数据，支持多线程加载和数据打乱
dataloaders = {
    # 训练集数据加载器，batch_size=32表示每批处理32张图片
    # shuffle=True表示每个epoch打乱数据顺序
    # num_workers=0表示不使用额外的子进程加载数据（Windows兼容）
    'train': DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0),
    # 验证集数据加载器，通常不需要打乱数据
    'val': DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=0)
}

# 获取训练集和验证集的大小
dataset_sizes = {
    'train': len(train_dataset),
    'val': len(val_dataset)
}

# 类别名称
class_names = ['cat', 'dog']
print(f"类别: {class_names}")
print(f"训练集大小: {dataset_sizes['train']}")
print(f"验证集大小: {dataset_sizes['val']}")

# 5. 加载经过预训练ResNet18模型
# 使用torchvision.models加载预训练的ResNet18模型
model = torchvision.models.resnet18(pretrained=True)

# 6. 冻结模型所有层的参数
# 遍历模型的所有参数，并将其requires_grad属性设置为False，这样在训练过程中不会更新这些参数
for param in model.parameters():
    param.requires_grad = False

# 7. 替换全连接层（输出2类）
# 获取全连接层的输入特征数
num_ftrs = model.fc.in_features
# 替换全连接层，使其输出2类（猫和狗）
model.fc = nn.Linear(num_ftrs, 2)

# 8. 配置GPU训练
# 将模型移动到指定设备（GPU或CPU）
model = model.to(device)

# 9. 编写训练循环（损失函数：CrossEntropyLoss，优化器：Adam）
# 定义损失函数，使用交叉熵损失函数（适用于分类问题）
criterion = nn.CrossEntropyLoss()

# 定义优化器，使用Adam优化器，优化替换后的全连接层的参数的参数（因为前面的层已经被冻结）
optimizer = optim.Adam(model.fc.parameters(), lr=0.0001)  # 11. 调整学习率（建议0.0001）

# 训练函数
def train_model(model, criterion, optimizer, num_epochs=20):
    """
    训练模型
    参数:
        model: 要训练的模型
        criterion: 损失函数
        optimizer: 优化器
        num_epochs: 训练轮数，默认为20轮
    返回:
        训练好的模型
    """
    # 记录最佳模型权重和最佳准确率
    best_model_wts = model.state_dict()
    best_acc = 0.0
    
    # 开始训练循环（20个epoch）
    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1}/{num_epochs}')
        print('-' * 10)
        
        # 每个epoch都有训练和验证阶段
        for phase in ['train', 'val']:
            if phase == 'train':
                # 设置模型为训练模式，启用Batch Normalization和Dropout
                model.train()
            else:
                # 设置模型为评估模式，禁用Batch Normalization和Dropout
                model.eval()
                
            # 初始化统计变量
            running_loss = 0.0
            running_corrects = 0
            
            # 遍历数据
            for inputs, labels in dataloaders[phase]:
                # 将数据移动到指定设备（GPU或CPU）
                inputs = inputs.to(device)
                labels = labels.to(device)
                
                # 清零梯度
                optimizer.zero_grad()
                
                # 前向传播
                #在验证阶段，梯度计算在 with 语句中被临时禁用，执行完代码块后会自动恢复为原来的状态（默认是启用状态）。
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    # 获取预测结果
                    #dim = 1 表示逐行取最大值所在的索引（0代表猫，1代表狗）
                    _, preds = torch.max(outputs, 1)
                    # 计算损失
                    loss = criterion(outputs, labels)
                    
                    # 反向传播和优化（仅在训练阶段）
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                
                # 统计损失和准确率
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            
            # 计算epoch的平均损失和准确率
            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            
            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
            
            # 深拷贝模型（在验证集上表现最好的模型）
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = model.state_dict()
                # 12. 保存最佳模型权重
                torch.save(model.state_dict(), 'best_model.pth')
        
        print()
    
    # 打印训练结果
    print(f'最佳验证集准确率: {best_acc:.4f}')
    
    # 加载最佳模型权重
    model.load_state_dict(best_model_wts)
    return model

# 10. 完整训练ResNet模型（20个epoch）
# 开始训练模型
model = train_model(model, criterion, optimizer, num_epochs=20)

print("训练完成！")